/*
 * Copyright (c) China Telecom Cloud Technology Co., Ltd. 2024-2025. All rights reserved.
 * Created by l00841256 on 25-3-31.
 */
#ifndef KVCACHESERVICE_H
#define KVCACHESERVICE_H

#include <grpcpp/grpcpp.h>

#include "ConfigLoader.h"
#include "KVStorageClient.h"
#include "Logger.h"
#include "MemoryPool.h"
#include "ThreadPool/ThreadPool.h"
#include <atomic>
#include <memory>
#include <mutex>
#include <shared_mutex>
#include <time.h>
#include <unordered_map>

#include "proto/kvcache.grpc.pb.h"
#include "proto/kvcache.pb.h"
#include "proto/kvctrl_api.grpc.pb.h"
#include "proto/kvctrl_api.pb.h"

#include "KVCachePublisher.h"


// 用于计算哈希的位配置
#define LEFT_BIT_NUM  6
#define RIGHT_BIT_NUM 2


//! 内存池大小
const uint64_t MEMPOOL_SIZE = ConfigLoader::get_config<uint64_t>(
    ConfigLoader::EnvVars::MEMPOOL_SIZE, DEFAULT_MEMPOOL_SIZE);


//! KV Cache存储根路径
const std::string MOUNT_POINT = ConfigLoader::get_config<std::string>(
    ConfigLoader::EnvVars::MOUNT_POINT, DEFAULT_MOUNT_POINT);


//! 块大小(即每个块包含多少token)
const uint64_t BLOCK_SIZE = ConfigLoader::get_config<uint64_t>(
    ConfigLoader::EnvVars::BLOCK_SIZE, DEFAULT_BLOCK_SIZE);


//! KV Cache读取线程数量
// TODO(yangxianpku： 2025.08.20)：考虑使用协程能
const int GET_THREADS = ConfigLoader::get_config<int>(
    ConfigLoader::EnvVars::GET_THREADS, DEFAULT_GET_THREADPOOL_SIZE);


//! KV Cache读取存储数量
// TODO(yangxianpku： 2025.08.20)：考虑使用协程能
const int PUT_THREADS = ConfigLoader::get_config<int>(
    ConfigLoader::EnvVars::PUT_THREADS, DEFAULT_PUT_THREADPOOL_SIZE);


/**
 *! @brief 实现gPRC服务端方法
 * 
 */
//! C++ 支持多继承
class KVCacheServiceImpl final : public kvcache::KVCacheService::Service,
                                 public KVCtrl::KVControlApi::Service {
private:
    std::atomic<bool> isReady = false;            // gPRC服务是否已经就绪

    std::unordered_set<size_t> perfixCacheSet;    // 前缀缓存集合(无重复)

    ThreadPool   getThreadPool;                   // KV Cache Get线程池
    ThreadPool   putThreadPool;                   // KV Cache Put线程池
    std::mutex   _poolsMutex;                     // 内存池修改的互斥锁

    /**
     *! @brief 查询_pools中是否已经有memSize块大小的内存池，如果没有就创建一个，
     *         返回pool指针并加入_pools中
     * 
     * @param memSize   内存池块大小
     * @param pool      获取到或创建的内存池指针
     * 
     * @return int32_t  成功返回0, 否则返回具体错误码
     */
    int32_t getMemoryPool(uint32_t memSize, MemoryPool **pool);



    std::unique_ptr<KVCachePublisher> _publisher;  // KV Cache事件信息发布者(Dynamo)
public:
    // 初始化列表中传入参数直接构造PUT与GET线程池
    explicit KVCacheServiceImpl()
        : getThreadPool(GET_THREADS), putThreadPool(PUT_THREADS) {}


    /**
     *! @brief  gRPC实现, 从存储中预读取KV Cache内容(先分析内存池中存在哪些，
     *          然后尝试从存储读取)
     * 
     * @param context 
     * @param request 
     * @param response 
     * @return grpc::Status 
     */
    grpc::Status PrefetchPrefixCache(grpc::ServerContext *context,
                        const kvcache::PrefetchPrefixCacheRequest *request,
                        kvcache::PrefetchPrefixCacheResponse *response);


    /**
     *! @brief 将请求中hashIds的对应的Segment设置为只读，需要时进行存储
     * 
     * @param context 
     * @param request 
     * @param response 
     * 
     * @return grpc::Status gRPC执行结果状态码
     */
    grpc::Status sealMemory(grpc::ServerContext *context,
                            const kvcache::sealMemoryRequest *request,
                            kvcache::sealMemoryResponse *response);


    /**
     *! @brief 通过块哈希值查询其信息: 内存块大小、共享文件路径、Segment属性
     * 
     * @param context 
     * @param request 
     * @param response 
     * @return grpc::Status 
     */
    grpc::Status QueryPrefixCache(grpc::ServerContext *context,
                    const kvcache::QueryPrefixCacheRequest *request,
                    kvcache::QueryPrefixCacheResponse *response
                ) override;


    /**
     *! @brief 通过块哈希值查询其信息: 内存块大小、共享文件路径、Segment属性，批量查询
     * 
     * @param context 
     * @param request 
     * @param response 
     * @return grpc::Status 
     */
    grpc::Status BatchQueryPrefixCache(
        grpc::ServerContext *context,
        const kvcache::BatchQueryPrefixCacheRequest *request,
        kvcache::BatchQueryPrefixCacheResponse *response) override;


    /**
     *! @brief 根据请求信息在内存池中分配Segments
     * 
     * @param context 
     * @param request 
     * @param response 
     * 
     * @return grpc::Status 
     */
    grpc::Status AllocateMemory(grpc::ServerContext *context,
                                const kvcache::AllocateMemoryRequest *request,
                                kvcache::AllocateMemoryResponse *response);


    /**
     *! @brief  获取指定内存池中所有已经Filled的Segment的哈希
     * 
     * @param context 
     * @param request 
     * @param response 
     * @return grpc::Status 
     */
    grpc::Status QueryKVDistribute(grpc::ServerContext *context,
                      const KVCtrl::QueryKVDistributeRequest *request,
                      KVCtrl::QueryKVDistributeReply *response);


    /**
     * @brief 
     * 
     * @param context 
     * @param request 
     * @param response 
     * @return grpc::Status 
     */
    grpc::Status Prefetch(grpc::ServerContext *context,
                          const KVCtrl::PrefetchRequest *request,
                          KVCtrl::PrefetchReply *response);


    /**
     *! @brief 根据请求参数获取对应的共享内存文件路径
     * 
     * @param context 
     * @param request 
     * @param response 
     * @return grpc::Status 
     */
    grpc::Status GetShmPath(grpc::ServerContext *context,
                            const kvcache::GetShmPathRequest *request,
                            kvcache::GetShmPathResponse *response);


    /**
     *! @brief 查看gRPC Server服务是否已经就绪
     * 
     * @param context 
     * @param request 
     * @param response 
     * @return grpc::Status 
     */
    grpc::Status CheckIfServiceReady(grpc::ServerContext *context,
                        const kvcache::CheckIfServiceReadyRequest *request,
                        kvcache::CheckIfServiceReadyResponse *response);

    /**
     *! @brief 根据gRPC远程请求参数删除对应MemoryPool
     * 
     * @param context  gPRC上下文(未使用)
     * @param request  内存池删除请求，带memsize参数
     * @param response 
     * 
     * @return grpc::Status 
     */
    grpc::Status CleanMemoryPool(grpc::ServerContext *context,
                                 const kvcache::CleanMemoryPoolRequest *request,
                                 kvcache::CleanMemoryPoolResponse *response);


    /**
     *! @brief 删除gRPC服务端，释放相关资源
    * 
    */
    ~KVCacheServiceImpl() override;


    /**
     *! @brief 初始化gRPC服务端环境, 包含如下内容：
     * 
     *  1. 创建get, put线程池(构造函数中万层)
     *  2. 创建一个默认块大小、总大小的内存池(共享， pinned)
     *  3. 初始化存储客户端
     *  4. 创建publisher
     *  5. 标识gPRC环境为就绪状态
     */
    void init();



    /**
     *! @brief 计算向量的哈希值 
     * 
     * @tparam T  泛型
     */
    template <typename T> 
    struct VectorHash {
        size_t operator()(const std::vector<T> &vec) const {
            std::hash<T> element_hasher;
            size_t seed = 0;
            for (const T &element : vec) {
                seed ^= element_hasher(element) + 0x9e3779b9 +
                        (seed << LEFT_BIT_NUM) + (seed >> RIGHT_BIT_NUM);
            }
            return seed;
        }
    };


    /**
     *! @brief 待哈希的块的任务
     * 
     */
    // TODO: 代码精简
    struct HashTask {
        //! 一个句子中某个未命中的块及以后所有块的串在一起做一个哈希
        // 假设有[blk0, blk1, ..., blki, ... blkn]
        // 假设[blk0, ..., blki] 在Mempool中命中， prefixCacheKey就是[blki+1, blkn]
        // 这些块的哈希的哈希。
        size_t      prefixCacheKey; 
        // 这一次未命中的哈希块中的某一哈希块
        BlockHashId blockHashId;
        Segment*    segment;

        HashTask(uint64_t prefixCacheKey, 
                BlockHashId blkHid, 
                Segment *s
            ) : prefixCacheKey(prefixCacheKey), 
            blockHashId(blkHid), segment(s) {}
    };


    /**
     *! @brief 序列状态管理
     * 
     */
    class SeqStatusManager {
    private:
        // TODO(yangxianpku： 2025.08.20)：minHash是什么含义？
        struct StatusEntry {
            time_t  lastUpdate;
            int64_t minHash;    //!  minHash 第一个未命中的哈希
        };

        // 是否已过期
        bool isEntryExpired(const StatusEntry &entry) const {
            time_t now = std::time(nullptr);
            return now - entry.lastUpdate > ENTRY_TTL;
        };

        time_t currentTimeForTesting = 0;
        mutable std::shared_mutex mutex;

        // key: prefixCacheKey
        std::unordered_map<size_t, StatusEntry> statusMap;
    public:
        // 记录保留时间
        static constexpr time_t ENTRY_TTL = 2;

        SeqStatusManager() {}

        ~SeqStatusManager() {}


        /**
         *! @brief statusMap没找到prefixCacheKey就要执行，找到了就不执行
         * 
         * @param task 
         * @return bool
         */
        bool shouldExecute(const HashTask &task) {
            const auto key = task.prefixCacheKey;
            {  // code block
                std::shared_lock readLock(mutex);  //! RAII
                auto it = statusMap.find(key);
                if (it == statusMap.end())
                    return true;   // not found
            }

            return false; // found
        }


        //! 从存储load该HashTask时从内存池中分配Segment失败或load文件失败
        void recordFailure(const HashTask &task) {
            const auto key = task.prefixCacheKey;
            {
                auto it = statusMap.find(key);
                if (it != statusMap.end())
                    return;
            }

            std::unique_lock writeLock(mutex);
            auto &entry      = statusMap[task.prefixCacheKey];

            entry.minHash    = std::min(entry.minHash, 
                                        task.blockHashId.hashId
                                    );
            entry.lastUpdate = std::time(nullptr);

            cleanupExpiredEntries(ENTRY_TTL);
        }


        void triggerCleanupForTesting() { 
            cleanupExpiredEntries(ENTRY_TTL); 
        }


        size_t getStatusMapSizeForTesting() {
            std::shared_lock lock(mutex);
            return statusMap.size();
        }


        /**
         *! @brief 清除掉statusMap中所有过期的条目
         * 
         * @param entryTtl 条目生存时间
         */
        void cleanupExpiredEntries(time_t entryTtl) {
            const auto now = std::time(nullptr);

            for (auto it = statusMap.begin(); it != statusMap.end();) {
                if (now - it->second.lastUpdate > entryTtl) {
                    it = statusMap.erase(it);
                } else {
                    ++it;
                }
            }
        }
    };  // SeqStatusManager


    //! key: memSize(内存块大小), value: MemoryPool
    // 可能存在多个内存池，每个内存池拥有特定的块大小
    std::unordered_map<uint32_t, std::unique_ptr<MemoryPool>> _pools;


    //! 存储客户端
    KVStorageClient &storageClient = KVStorageClient::getInstance();


    SeqStatusManager statusManager;


private:
    /**
     *! @brief 根据segmentId查找所在的Memory Pool，如果没找到返回nullptr
     * 
     * @param hashId    需要查询的hashId，一般为seq的首个hashId
     * @param pool      查找到的Memory Pool指针，找不到时为nullptr
     * 
     * @return int32_t  查找成功为0，否则返回具体错误码
     */
    // TODO(yangxianpku: 2025.08.20): 这个函数可能比较耗时
    int32_t getMemoryPoolByHashId(int64_t hashId, MemoryPool **pool);


    /**
     *! @brief 构建从存储读取KV Cache数据到内存池的异步future(内存池分配Segment+存储Load)
     * 
     * @param taskFutures 
     * @param pool 
     * @param getCacheTask 待从存储捞取的KV Cache块的HashTask
     * @param tpNum 
     * @param cacheInfo 
     */
    void buildPrefetchPrefixCacheTaskFuture(
                std::vector<std::future<int32_t>> &taskFutures, 
                MemoryPool *pool,
                HashTask getCacheTask, 
                int32_t tpNum, 
                CacheInfo cacheInfo
            );
    


    /**
     *! @brief 查询指定请求中包含的序列(序列包含块)有哪些块的哈希不在内存池中
     * 
     * @param request       gRPC客户端请求
     * @param pool          内存池
     * @param getCacheTasks [out] 返回内存池中没有的哈希块的HashTask
     * @param allHashIds    [out] 返回该请求包含的所有块哈希
     */
    void processPrefetchPrefixCacheRequest(
                                const kvcache::PrefetchPrefixCacheRequest *request, 
                                MemoryPool *pool,
                                std::vector<HashTask> &getCacheTasks, 
                                std::vector<int64_t>  &allHashIds
                            );
    

    /**
     *! @brief 提交写入任务到写入线程池: 将对应的Segment写入存储系统
     * 
     * @param pool 
     * @param blockHashId 
     * @param tpNum 
     */
    void submitWriteSegmentTask(MemoryPool *pool, 
                                BlockHashId blockHashId,
                                int32_t tpNum
                            );


    /**
     * @brief 
     * 
     * @param taskFutures 
     * @param pool 
     * @param getCacheTask 
     * @param tpNum 
     */
    void buildPrefetchTaskFuture(std::vector<std::future<int32_t>> &taskFutures,
                                MemoryPool *pool, 
                                HashTask getCacheTask,
                                int32_t tpNum
                            );
    

    /**
     * @brief 
     * 
     * @param request 
     * @param pool 
     * @param getCacheTasks 
     * @return int32_t 
     */
    int32_t processPrefetchRequest(const KVCtrl::PrefetchRequest *request,
                                MemoryPool *pool,
                                std::vector<HashTask> &getCacheTasks
                            );
};
#endif